import torch
import traceback, pdb
from infer_uvr5 import _audio_pre_, _audio_pre_new
from lib.audio import load_audio
import MDXNet
import os
from config import Config

config = Config()
weight_root = "weights"
weight_uvr5_root = "uvr5_weights"
index_root = "logs"
names = []

for name in os.listdir(weight_root):
    if name.endswith(".pth"):
        names.append(name)
index_paths = []
for root, dirs, files in os.walk(index_root, topdown=False):
    for name in files:
        if name.endswith(".index") and "trained" not in name:
            index_paths.append("%s/%s" % (root, name))
uvr5_names = []
for name in os.listdir(weight_uvr5_root):
    if name.endswith(".pth") or "onnx" in name:
        uvr5_names.append(name.replace(".pth", ""))
now_dir = os.getcwd()
tmp = os.path.join(now_dir, "TEMP")


def uvr(model_name, inp_root, save_root_vocal, paths, save_root_ins, agg, format0):
    infos = []
    try:
        inp_root = inp_root.strip(" ").strip('"').strip("\n").strip('"').strip(" ")
        save_root_vocal = (
            save_root_vocal.strip(" ").strip('"').strip("\n").strip('"').strip(" ")
        )
        save_root_ins = (
            save_root_ins.strip(" ").strip('"').strip("\n").strip('"').strip(" ")
        )
        if model_name == "onnx_dereverb_By_FoxJoy":
            from MDXNet import MDXNetDereverb

            pre_fun = MDXNetDereverb(15)
        else:
            func = _audio_pre_ if "DeEcho" not in model_name else _audio_pre_new
            pre_fun = func(
                agg=int(agg),
                model_path=os.path.join(weight_uvr5_root, model_name + ".pth"),
                device=config.device,
                is_half=config.is_half,
            )
        if inp_root != "":
            paths = [os.path.join(inp_root, name) for name in os.listdir(inp_root)]
        else:
            paths = [path.name for path in paths]
        for path in paths:
            inp_path = os.path.join(inp_root, path)
            need_reformat = 1
            done = 0
            try:
                info = ffmpeg.probe(inp_path, cmd="ffprobe")
                if (
                        info["streams"][0]["channels"] == 2
                        and info["streams"][0]["sample_rate"] == "44100"
                ):
                    need_reformat = 0
                    pre_fun._path_audio_(
                        inp_path, save_root_ins, save_root_vocal, format0
                    )
                    done = 1
            except:
                need_reformat = 1
                traceback.print_exc()
            if need_reformat == 1:
                tmp_path = "%s/%s.reformatted.wav" % (tmp, os.path.basename(inp_path))
                os.system(
                    "ffmpeg -i %s -vn -acodec pcm_s16le -ac 2 -ar 44100 %s -y"
                    % (inp_path, tmp_path)
                )
                inp_path = tmp_path
            try:
                if done == 0:
                    pre_fun._path_audio_(
                        inp_path, save_root_ins, save_root_vocal, format0
                    )
                infos.append("%s->Success" % (os.path.basename(inp_path)))
                yield "\n".join(infos)
            except:
                infos.append(
                    "%s->%s" % (os.path.basename(inp_path), traceback.format_exc())
                )
                yield "\n".join(infos)
    except:
        infos.append(traceback.format_exc())
        yield "\n".join(infos)
    finally:
        try:
            if model_name == "onnx_dereverb_By_FoxJoy":
                del pre_fun.pred.model
                del pre_fun.pred.model_
            else:
                del pre_fun.model
                del pre_fun
        except:
            traceback.print_exc()
        print("clean_empty_cache")
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    yield "\n".join(infos)


# 设置文件路径

def rename(path, part):
    # 遍历目录中的所有文件
    for filename in os.listdir(path):
        # 检查文件是否在指定的目录下
        if os.path.isfile(os.path.join(path, filename)):
            # 分割文件名和扩展名
            file_root, file_ext = os.path.splitext(filename)
            print(file_ext, filename, file_root)
            if file_ext not in ['.flac', '.m4a', '.wav']: continue
            # 假设我们要保留文件名中的".m4a"之前的部分
            if part == 'vocal':
                new_root = filename.replace('instrument', 'vocal')
            else:
                new_root = filename.replace('vocal', 'instrument')

            # 构建完整的原文件路径和新文件路径
            old_file = os.path.join(path, filename)
            new_file = os.path.join(path, new_root)
            # 重命名文件
            print('old new', old_file, new_file)
            os.rename(old_file, new_file)
            return new_file

def run_sample():
    model_name = 'HP3_all_vocals'
    inp_root = '/root/autodl-tmp/newTrain'
    save_root_vocal = '/root/autodl-tmp/vocal'
    save_root_ins = '/root/autodl-tmp/inst'
    agg = 10
    format0 = 'flac'
    print('run')
    res = uvr(model_name, inp_root, save_root_vocal, [], save_root_ins, agg, format0)
    for r in res:
        print(r)
    rename(save_root_vocal, 'vocal')
    rename(save_root_ins, 'instrument')


def run_model(inp_root, save_root_vocal, save_root_ins, format0,dorename = False):
    model_name = 'HP3_all_vocals'
    # inp_root  = '/root/autodl-tmp/sample_train'
    # save_root_vocal = '/root/autodl-tmp/vocal'
    # save_root_ins= '/root/autodl-tmp/inst'
    agg = 10
    # format0 = 'flac'
    print('run')
    print('运行uvr',inp_root,save_root_vocal,save_root_ins)
    res = uvr(model_name, inp_root, save_root_vocal, [], save_root_ins, agg, format0)
    for r in res:
        print(r)
    if dorename:
        rename(save_root_vocal, 'vocal')
        rename(save_root_ins, 'instrument')

# inp_root='/root/autodl-tmp/Retrieval-based-Voice-Conversion-WebUI/bksong'
# file_name_without_extension ='huilongdajie'
# save_root_vocal = f'/root/autodl-tmp/Retrieval-based-Voice-Conversion-WebUI/split_{file_name_without_extension}/vocal'
# save_root_ins= f'/root/autodl-tmp/Retrieval-based-Voice-Conversion-WebUI/split_{file_name_without_extension}/inst'
# format0 = 'mp3'
# inp_root = '/root/autodl-tmp/huazhou'
# save_root_vocal = '/root/autodl-tmp/huazhou_vocal'
# save_root_ins = '/root/autodl-tmp/huazhou_instru'
# # #run_sample()
# run_model(inp_root,
#           save_root_vocal,
#           save_root_ins,
#           format0)
